import argparse
import os
import sys
import traceback
import zipfile

import requests
from tqdm import tqdm

ROOT_URL = 'http://{}:{}/download/{}/'

DB_ROOT = f'{sys.path[0]}/db'
MODEL_ROOT = f'{sys.path[0]}/model'

DB = 'db.zip'
MODEL = 'model.zip'


def __get_db(root_url, root_path, file_name):
    url = str(root_url + file_name)
    resp = requests.get(url, stream=True)
    if resp.status_code != 200:
        raise RuntimeError("download error")
    try:
        total = round(int(resp.headers['Content-Length']) / 1024 / 1024)
        with open(os.path.join(root_path, file_name), 'wb') as f:
            for data in tqdm(iterable=resp.iter_content(1024 * 1024), ncols=100,
                             desc='Downloading...', total=total, unit='MB'):
                f.write(data)
    except IOError or RuntimeError:
        traceback.print_exc(limit=10, file=sys.stdout)


def __update(ip, port, topic, root_path, file_name):
    if not os.path.exists(root_path):
        os.mkdir(root_path)
    root_url = ROOT_URL.format(ip, port, topic)
    __get_db(root_url, root_path, file_name)
    zip_file = os.path.join(root_path, file_name)
    with zipfile.ZipFile(zip_file, 'r') as z:
        for f in tqdm(z.namelist()[:1000], ncols=100, desc="Extracting...", unit='files'):
            z.extract(f, path=root_path)
    z.close()
    os.remove(zip_file)


def check_file(ds_ip, bed_port, is_test=False, force_update=False):
    update_db(ds_ip, bed_port, is_test)
    if not os.path.exists(MODEL_ROOT) or force_update:
        update_model(ds_ip, bed_port, is_test)


def update_db(ds_ip, bed_port, is_test=False):
    typ = 'tdb' if is_test else 'db'
    __update(ds_ip, bed_port, typ, DB_ROOT, DB)


def update_model(ds_ip, bed_port, is_test=False):
    __update(ds_ip, bed_port, 'model', MODEL_ROOT, MODEL)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('ds_ip', type=str,
                        help='IP address to download.')
    parser.add_argument('--is_test', action="store_true",
                        help='Whether to download test file or not.')
    parser.add_argument('--force_update', action="store_true",
                        help='Whether to force update (update model).')
    args = parser.parse_args(sys.argv[1:])
    exec("check_file('{}', '{}', '{}')".format(
        str(args.ds_ip), args.is_test, args.force_update))
